# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/5/19 15:57
# @Function:
from typing import Tuple, Union
import os
import numpy as np
import xarray as xr
from xarray import Dataset, DataArray
# from rioxarray.rioxarray import DataArray as RioDataArray



def compute_relative_humidity(tas: xr.DataArray, tdps: xr.DataArray) -> xr.DataArray:
    """
    根据气温和露点温度计算相对湿度
    tas 和 tdps 单位应为摄氏度
    """

    def saturation_vapor_pressure(temp_c):
        # Magnus formula for vapor pressure in hPa (mb)
        return 6.112 * np.exp((17.67 * temp_c) / (temp_c + 243.5))

    es = saturation_vapor_pressure(tas)  # 饱和水汽压
    e = saturation_vapor_pressure(tdps)  # 实际水汽压

    rh = (e / es) * 100
    rh = rh.clip(0, 100)  # 限制在 0~100%
    return rh.astype(np.float32)


def make_humidity_dataset(temp_file: str, dewpoint_file: str, output_file: str):
    # 加载数据集
    ds_temp = xr.open_dataset(temp_file)
    ds_dew = xr.open_dataset(dewpoint_file)

    # todo 未作数据校验
    # 时间维度对齐
    # common_time = ds_temp.valid_time.intersection(ds_dew.valid_time)
    # ds_temp = ds_temp.sel(valid_time=common_time)
    # ds_dew = ds_dew.sel(valid_time=common_time)

    # 提取变量（假设变量名为 t2m 和 d2m）
    tas = ds_temp[list(ds_temp.data_vars.keys())[0]] - 273.15  # 转换为摄氏度
    tdps = ds_dew[list(ds_dew.data_vars.keys())[0]] - 273.15  # 转换为摄氏度

    # 计算相对湿度
    relative_humidity = compute_relative_humidity(tas, tdps)

    # 创建新的 Dataset
    ds_rh = xr.Dataset(
        {
            'relative_humidity': relative_humidity,
        },
        coords={
            'valid_time': relative_humidity.valid_time,
            'latitude': relative_humidity.latitude,
            'longitude': relative_humidity.longitude,
        },
        attrs={
            'description': 'Relative Humidity calculated from dew point and temperature',
            'units': '%',
            'source_temperature': os.path.basename(temp_file),
            'source_dew_point': os.path.basename(dewpoint_file),
        }
    )

    # 保存为 NetCDF 文件
    ds_rh.to_netcdf(output_file)
    print(f"已保存湿度数据至 {output_file}")



def load_datasets(evaporation_file: str, temperature_file: str) -> Tuple[Dataset, Dataset]:
    """加载蒸发和气温数据集"""
    evaporation_dataset = xr.open_dataset(evaporation_file)
    temperature_dataset = xr.open_dataset(temperature_file)
    return evaporation_dataset, temperature_dataset

def align_time_dimension(evaporation_dataset: Dataset, temperature_dataset: Dataset) -> Dataset:
    """确保两个数据集的时间维度一致"""
    return evaporation_dataset.sel(valid_time=temperature_dataset['valid_time'])

def create_temperature_mask(temperature_dataset: Dataset, threshold_k: float) -> DataArray:
    """创建气温大于指定阈值（单位：K）的掩膜"""
    temp_var = list(temperature_dataset.data_vars.keys())[0]
    return temperature_dataset[temp_var] > threshold_k

def create_humidity_mask(humidity_dataset: Dataset, threshold_percent: float) -> DataArray:
    """创建湿度小于指定阈值（百分比）的掩膜"""
    hum_var = list(humidity_dataset.data_vars.keys())[0]
    return humidity_dataset[hum_var] < threshold_percent


def calculate_evaporation_during_specific_period(evaporation_dataset: Dataset, mask: DataArray) -> DataArray:
    """计算空调降温期间的蒸散发量"""
    e_var = list(evaporation_dataset.data_vars.keys())[0]
    e_wt = evaporation_dataset[e_var].where(mask, drop=True)
    annual_e_wt = e_wt.resample(valid_time='1YE').sum(dim='valid_time')
    return annual_e_wt

def compute_hours_to_days(mask: DataArray) -> np.ndarray:
    """统计每个格点气温大于阈值的小时数"""
    num_hours_gt_26 = mask.sum(dim='valid_time').values
    num_days_gt_26 = np.ceil(num_hours_gt_26 / 24)
    num_days_gt_26[num_days_gt_26 == 0] = 1  # 避免零天情况
    return num_days_gt_26

def save_to_netcdf_and_tiff(data_array: DataArray, output_dir: str, year: int,
                            nodata: Union[int, float] = -9999, prefix: str = "Ewt") -> Tuple[str, str]:
    """保存为 NetCDF 和 GeoTIFF 格式"""
    # NetCDF 输出
    nc_file = os.path.join(output_dir, f'{prefix}_{year}.nc')
    data_array.to_netcdf(nc_file)

    # GeoTIFF 输出
    tif_file = os.path.join(output_dir, f'{prefix}_{year}.tif')

    data_rio = data_array.rio.write_crs("EPSG:4326")
    data_rio = data_rio.rio.write_nodata(nodata)
    data_rio = data_rio.astype(np.float32)
    data_rio.rio.to_raster(tif_file)

    print(f"数据已保存到 {nc_file}")
    print(f"数据已保存到 {tif_file}")

    return nc_file, tif_file

def cool_main(input_dir, year, output_dir=None, temperature_file_name='2023_temp_hour.nc',
         evaporation_file_name='Evaporation_hour_2023.nc', temp_threshold_c=26.0, nodata=-9999):
    """
    主流程函数
    :param input_dir: 输入数据目录
    :param year: 年份
    :param output_dir: 输出目录，若不传则默认为 input_dir/output
    :param temperature_file_name: 气温文件名
    :param evaporation_file_name: 蒸发文件名
    :param temp_threshold_c: 温度阈值（摄氏度）
    :param nodata: 设置无效值
    """
    if output_dir is None:
        output_dir = os.path.join(input_dir, 'output')
        os.makedirs(output_dir, exist_ok=True)

    temperature_file = os.path.join(input_dir, temperature_file_name)
    evaporation_file = os.path.join(input_dir, evaporation_file_name)

    # 加载数据
    evap_ds, temp_ds = load_datasets(evaporation_file, temperature_file)

    # 时间对齐
    evap_ds = align_time_dimension(evap_ds, temp_ds)

    # 创建掩膜
    threshold_k = temp_threshold_c + 273.15
    mask = create_temperature_mask(temp_ds, threshold_k)

    # 计算蒸散量
    annual_e_wt = calculate_evaporation_during_specific_period(evap_ds, mask)

    # 保存EWT结果
    save_to_netcdf_and_tiff(annual_e_wt, output_dir, year, nodata=nodata, prefix="Ewt")

    # 统计高温天数
    num_days_gt_26 = compute_hours_to_days(mask)

    # 构建DataArray并保存为TIFF
    da_days = xr.DataArray(
        num_days_gt_26,
        dims=['latitude', 'longitude'],
        coords={'latitude': annual_e_wt.latitude, 'longitude': annual_e_wt.longitude}
    )
    save_to_netcdf_and_tiff(da_days, output_dir, year, nodata=nodata, prefix="ETT-D")

    # 关闭数据集
    evap_ds.close()
    temp_ds.close()

def dry_main(input_dir, year, output_dir=None,
         humidity_file_name='2023_temp_hour.nc',
         evaporation_file_name='Evaporation_hour_2023.nc',
         humidity_threshold=40, nodata=-9999):
    """
    主流程函数
    :param input_dir: 输入数据目录
    :param year: 年份
    :param output_dir: 输出目录，若不传则默认为 input_dir/output
    :param humidity_file_name: 湿度文件名
    :param evaporation_file_name: 蒸发文件名
    :param humidity_threshold: 湿度阈值（百分比）
    :param nodata: 设置无效值
    """
    if output_dir is None:
        output_dir = os.path.join(input_dir, 'output')
        os.makedirs(output_dir, exist_ok=True)

    humidity_file = os.path.join(input_dir, humidity_file_name)
    evaporation_file = os.path.join(input_dir, evaporation_file_name)

    # 加载数据
    evap_ds, humid_ds = load_datasets(evaporation_file, humidity_file)

    # 时间对齐
    evap_ds = align_time_dimension(evap_ds, humid_ds)

    # 创建掩膜
    mask = create_humidity_mask(humid_ds, humidity_threshold)

    # 计算蒸散量
    annual_e_wh = calculate_evaporation_during_specific_period(evap_ds, mask)

    # 保存EWH结果
    save_to_netcdf_and_tiff(annual_e_wh, output_dir, year, nodata=nodata, prefix="Ewh")

    # 统计低湿度天数
    num_days_gh_40 = compute_hours_to_days(mask)

    # 构建DataArray并保存为TIFF
    da_days = xr.DataArray(
        num_days_gh_40,
        dims=['latitude', 'longitude'],
        coords={'latitude': annual_e_wh.latitude, 'longitude': annual_e_wh.longitude}
    )
    save_to_netcdf_and_tiff(da_days, output_dir, year, nodata=nodata, prefix="Ewh-D")

    # 关闭数据集
    evap_ds.close()
    humid_ds.close()

# 示例调用
if __name__ == '__main__':
    # input_folder = r'G:\GEP_data\era5\2023era5data\python_download\reanalysis-era5-single-levels\china'
    # cool_main(input_dir = input_folder, year = 2023, temperature_file_name='single-levels_2m_temperature_2023.nc',
    #           evaporation_file_name='single-levels_evaporation_2023.nc')

    # dry_main(input_dir = input_folder, year = 2023, humidity_file_name='2023_temp_hour.nc',
    #          evaporation_file_name='single-levels_evaporation_2023.nc')

    # input_dir = r'G:\GEP_data\era5\2024era5data'


    input_folder = r'G:\GEP_data\era5\2023era5data\python_download\reanalysis-era5-single-levels\china'
    year = 2024
    temp_name = f"single-levels_2m_temperature_{year}.nc"
    dewpoint_name = f"single-levels_2m_dewpoint_temperature_{year}.nc"
    evaporation_name = f"single-levels_evaporation_{year}.nc"
    humidity_name = f"2m_humidity_{year}.nc"

    # cool_main(input_dir = input_folder, year=year, temperature_file_name=temp_name,
    #           evaporation_file_name=evaporation_name)
    #
    # make_humidity_dataset(
    #     temp_file = os.path.join(input_folder, temp_name),
    #     dewpoint_file= os.path.join(input_folder, dewpoint_name),
    #     output_file = os.path.join(input_folder, humidity_name)
    #     )

    dry_main(input_dir = input_folder, year=year, humidity_file_name=humidity_name,
             evaporation_file_name=evaporation_name)
